import os
import tensorflow as tf
from tensorflow.keras import layers
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split

import sys
sys.path.append('..')


class bpnn():
    """
    BP神经网络
    输入    中间层    中间层    输出层
    7        5        4        1
    """

    tf.random.set_seed(100)
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

    def __init__(self):
        self.model = tf.keras.models.Sequential()
        # 输入7维特征向量
        # 5个神经元的和4个神经元的中间层
        self.model.add(layers.Dense(5, activation='relu', input_shape=(7, )))
        self.model.add(layers.Dense(4, activation='relu'))
        # 输出层1个神经元
        self.model.add(layers.Dense(1, activation='sigmoid'))
        # 配置模型学习流程
        self.model.compile(optimizer='adam',
                           loss='binary_crossentropy',
                           metrics=['accuracy'])
        self.model.summary()

    # 训练
    def fit(self, train_x, train_y, epochs=100, batch_size=100, validation_data=None):
        history = self.model.fit(train_x, train_y, epochs=epochs, batch_size=batch_size, validation_data=validation_data)
        return history

    # 评估
    def evaluate(self, test_x, test_y, batch_size=32):
        self.model.evaluate(test_x, test_y, batch_size=batch_size)

    # 预测
    def predict(self, x):
        return self.model.predict(x)

    # 保存模型
    def save(self):
        self.model.save('../saved_model/bpnn_model_input7.h5')


class simple_dnn():
    """
    简单 DNN 神经网络
    输入    中间层    中间层    中间层    输出层
    7        12        6         3        4
    """

    tf.random.set_seed(100)
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

    def __init__(self):
        self.model = tf.keras.models.Sequential()

        self.model.add(layers.Dense(12, activation='relu', input_shape=(7, )))
        self.model.add(layers.Dense(6, activation='relu'))
        self.model.add(layers.Dense(3, activation='relu'))
        # 输出层2个神经元
        self.model.add(layers.Dense(4, activation='softmax'))
        # 配置模型学习流程
        self.model.compile(optimizer='adam',
                           loss='binary_crossentropy',
                           metrics=['accuracy'])
        self.model.summary()

    # 训练
    def fit(self, train_x, train_y, epochs=500, batch_size=100, validation_data=None):
        history = self.model.fit(train_x, train_y, epochs=epochs, batch_size=batch_size, validation_data=validation_data)
        return history

    # 评估
    def evaluate(self, test_x, test_y, batch_size=32):
        self.model.evaluate(test_x, test_y, batch_size=batch_size)

    # 预测
    def predict(self, x):
        pre = self.model.predict(x)
        return np.argmax(pre[0])

    # 保存模型
    def save(self):
        self.model.save('../saved_model/dnn_model_input7.h5')


if __name__ == '__main__':
    pass
